from ultralytics import YOLO
import torch

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Load a model
model = YOLO("yolov8n.yaml")  # .load('yolov8n.pt')#'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
model.to(device)
# Train the model
model.train(data="data.yaml", epochs=50, imgsz=256, batch=5, ch=1, close_mosaic=50)
